// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/exec/aggregate.h>
#include <pollux/expression/function_signature.h>

namespace kumo::pollux::exec {
    class CompanionSignatures {
    public:
        // Return signatures of the partial companion function for the original
        // aggregate function of `signatures`.
        static std::vector<AggregateFunctionSignaturePtr> partialFunctionSignatures(
            const std::vector<AggregateFunctionSignaturePtr> &signatures);

        static std::string partialFunctionName(const std::string &name);

        static AggregateFunctionSignaturePtr mergeFunctionSignature(
            const AggregateFunctionSignaturePtr &signature);

        // Return signatures of the merge companion function for the original
        // aggregate function of `signatures`.
        static std::vector<AggregateFunctionSignaturePtr> mergeFunctionSignatures(
            const std::vector<AggregateFunctionSignaturePtr> &signatures);

        static std::string mergeFunctionName(const std::string &name);

        // Return true if there are multiple signatures that have the same
        // intermediate type.
        static bool hasSameIntermediateTypesAcrossSignatures(
            const std::vector<AggregateFunctionSignaturePtr> &signatures);

        static AggregateFunctionSignaturePtr mergeExtractFunctionSignature(
            const AggregateFunctionSignaturePtr &signature);

        static std::vector<AggregateFunctionSignaturePtr>
        mergeExtractFunctionSignatures(
            const std::vector<AggregateFunctionSignaturePtr> &signatures);

        static std::string mergeExtractFunctionNameWithSuffix(
            const std::string &name,
            const TypeSignature &resultType);

        static std::string mergeExtractFunctionName(const std::string &name);

        // Return signature of the extract companion function for the original
        // aggregate function of `signature`.
        static FunctionSignaturePtr extractFunctionSignature(
            const AggregateFunctionSignaturePtr &signature);

        static std::string extractFunctionNameWithSuffix(
            const std::string &name,
            const TypeSignature &resultType);

        // Return signatures of the extract companion function for the original
        // aggregate function of `signatures`.
        static std::vector<FunctionSignaturePtr> extractFunctionSignatures(
            const std::vector<AggregateFunctionSignaturePtr> &signatures);

        static std::string extractFunctionName(const std::string &name);

        static std::
        unordered_map<TypeSignature, std::vector<AggregateFunctionSignaturePtr> >
        groupSignaturesByReturnType(
            const std::vector<AggregateFunctionSignaturePtr> &signatures);

    private:
        static TypeSignature normalizeTypeImpl(
            const TypeSignature &type,
            const melon::F14FastMap<std::string, SignatureVariable> &allVariables,
            melon::F14FastMap<std::string, std::string> &renamedVariables);

        // Rename variables used in `type` as T0, T1, ..., Tn, with the ordering of
        // variables being visited in pre-order traversal.
        static TypeSignature normalizeType(
            const TypeSignature &type,
            const melon::F14FastMap<std::string, SignatureVariable> &allVariables);

        // Process signatures of distinct pairs of intermediate and result types and
        // return a vector of the processed signatures.
        template<class T>
        static std::vector<T> processSignaturesOfDistinctTypes(
            const std::vector<AggregateFunctionSignaturePtr> &signatures,
            const std::function<T(const AggregateFunctionSignaturePtr &)> &
            processSignature) {
            std::unordered_set<std::pair<TypeSignature, TypeSignature> >
                    distinctIntermediateAndResultTypes;
            std::vector<T> processedSignatures;

            for (const auto &signature: signatures) {
                auto normalizedIntermediateType =
                        normalizeType(signature->intermediateType(), signature->variables());
                auto normalizedReturnType =
                        normalizeType(signature->returnType(), signature->variables());
                if (distinctIntermediateAndResultTypes.count(std::make_pair(
                    normalizedIntermediateType, normalizedReturnType))) {
                    continue;
                }

                auto processedSignature = processSignature(signature);
                if (processedSignature) {
                    processedSignatures.push_back(std::move(processedSignature));
                    // Only remember the intermediate and result types and skip subsequent
                    // signatures of the same types if we have already successfully
                    // processed one of them. There may be multiple signatures of the same
                    // intermediate type, some can be processed successfully while some
                    // cannot. We only need one processed signature for each pair of
                    // intermediate and result types.
                    distinctIntermediateAndResultTypes.emplace(
                        normalizedIntermediateType, normalizedReturnType);
                }
            }
            return processedSignatures;
        }

        // Process signatures of distinct intermediate types and return a vector of
        // the processed signatures.
        static std::vector<AggregateFunctionSignaturePtr>
        processSignaturesOfDistinctIntermediateTypes(
            const std::vector<AggregateFunctionSignaturePtr> &signatures,
            const std::function<AggregateFunctionSignaturePtr(
                const AggregateFunctionSignaturePtr &)> &processSignature) {
            std::unordered_set<TypeSignature> distinctIntermediateTypes;
            std::vector<AggregateFunctionSignaturePtr> processedSignatures;

            for (const auto &signature: signatures) {
                auto normalizedIntermediateType =
                        normalizeType(signature->intermediateType(), signature->variables());
                if (distinctIntermediateTypes.count(normalizedIntermediateType)) {
                    continue;
                }

                auto processedSignature = processSignature(signature);
                if (processedSignature) {
                    processedSignatures.push_back(std::move(processedSignature));
                    // Only remember the intermediate type and skip subsequent signatures of
                    // the same type if we have already successfully processed one of them.
                    // There may be multiple signatures of the same intermediate type, some
                    // can be processed successfully while some cannot. We only need one
                    // processed signature for each intermediate type.
                    distinctIntermediateTypes.emplace(normalizedIntermediateType);
                }
            }
            return processedSignatures;
        }
    };

    // Return type variables used in `types`.
    melon::F14FastMap<std::string, SignatureVariable> usedTypeVariables(
        const std::vector<TypeSignature> &types,
        const melon::F14FastMap<std::string, SignatureVariable> &allVariables);

    // Return true if the result type of `signature` can be resolved from a concrete
    // intermediate type of it.
    bool isResultTypeResolvableGivenIntermediateType(
        const AggregateFunctionSignaturePtr &signature);

    // Return a string that is preorder traversal of `type`. For example, for
    // row(bigint, array(double)), return a string "row_bigint_array_double".
    std::string toSuffixString(const TypeSignature &type);
} // namespace kumo::pollux::exec
